from PyQt5.QtWidgets import (QHBoxLayout, QVBoxLayout, QPushButton,
                             QFileDialog, QMessageBox, QWidget,
                             QDialog)
from PyQt5.QtCore import Qt, QRect, QPoint, pyqtSignal,QSize
from PyQt5.QtGui import QPixmap, QPainter, QPen, QColor, QGuiApplication, QCursor, QIcon

import sys

class ScreenshotDialog(QDialog):
    def __init__(self):
        super().__init__()
        self.initUI()

        # 截图相关变量
        self.selection_start = None
        self.selection_end = None
        self.is_selecting = False
        self.is_moving = False
        self.move_offset = QPoint()
        self.selected_region = None

        # 获取屏幕截图
        self.capture_screen()

    def initUI(self):
        self.setWindowTitle("高级截图工具")
        self.setWindowFlags(Qt.WindowStaysOnTopHint)

        # 控制面板
        self.control_panel = QWidget(self)
        self.control_panel.setFixedWidth(120)
        self.control_panel.hide()  # 初始隐藏控制面板

        self.save_btn = QPushButton(self)
        self.save_btn.setIcon(QIcon(QPixmap("pic/ok.png")))
        self.save_btn.setIconSize(QSize(30,30))
        self.save_btn.setFixedSize(QSize(40,40))
        self.save_btn.setFlat(True)
        self.save_btn.clicked.connect(self.save_screenshot)

        # 控制面板布局
        panel_layout = QHBoxLayout()
        panel_layout.setContentsMargins(0,0,0,0)
        panel_layout.addStretch()
        panel_layout.addWidget(self.save_btn)

        self.control_panel.setLayout(panel_layout)

        # 主布局
        main_layout = QVBoxLayout()
        main_layout.addWidget(self.control_panel)
        main_layout.setAlignment(Qt.AlignTop)
        self.setLayout(main_layout)

        # 初始位置
        self.move(50, 50)
        self.show()

    def capture_screen(self):
        """获取屏幕截图并进入选择模式"""
        screen = QGuiApplication.primaryScreen()
        self.screenshot = screen.grabWindow(0)

        # 重置选择状态
        self.selection_start = None
        self.selection_end = None
        self.is_selecting = False
        self.is_moving = False
        self.selected_region = None
        self.control_panel.hide()

        # 设置窗口为全屏
        self.setGeometry(0, 0, self.screenshot.width(), self.screenshot.height())
        self.showFullScreen()
        self.setCursor(Qt.CrossCursor)
        self.update()

    def paintEvent(self, event):
        painter = QPainter(self)

        # 绘制屏幕截图
        painter.drawPixmap(0, 0, self.screenshot)

        # 绘制半透明蒙版
        painter.setBrush(QColor(0, 0, 0, 160))
        painter.setPen(Qt.NoPen)
        painter.drawRect(0, 0, self.width(), self.height())

        # 绘制选择区域
        if self.selection_start and self.selection_end:
            rect = QRect(self.selection_start, self.selection_end).normalized()

            # 绘制原始截图区域
            painter.drawPixmap(rect, self.screenshot, rect)

            # 绘制选择框
            painter.setPen(QPen(QColor(255, 100, 100), 2, Qt.DashLine))
            painter.setBrush(Qt.NoBrush)
            painter.drawRect(rect)

            # 显示尺寸信息
            painter.setPen(QPen(Qt.white))
            info_text = f"{rect.width()} × {rect.height()}  (X:{rect.x()}, Y:{rect.y()})"
            painter.drawText(rect.bottomRight() + QPoint(10, 20), info_text)

            # 绘制调整手柄
            if not self.is_selecting:
                self.draw_resize_handles(painter, rect)

                # 显示控制面板在选区底部
                self.show_control_panel(rect)

    def show_control_panel(self, selection_rect):
        """在选区底部显示控制面板"""
        if not self.control_panel.isVisible():
            self.control_panel.show()

        # 计算控制面板位置
        panel_width = 180  # 控制面板宽度
        panel_height = 30  # 控制面板高度

        # 确保控制面板不会超出屏幕
        panel_x = selection_rect.x() + (selection_rect.width() - panel_width) // 2
        panel_y = selection_rect.bottom() + 5

        # 如果底部空间不足，显示在选区上方
        if panel_y + panel_height > self.height():
            panel_y = selection_rect.top() - panel_height - 10

        # 如果右侧空间不足，向左偏移
        if panel_x + panel_width > self.width():
            panel_x = self.width() - panel_width - 10

        # 如果左侧空间不足，向右偏移
        if panel_x < 0:
            panel_x = 10

        self.control_panel.setGeometry(panel_x, panel_y, panel_width, panel_height)

    def draw_resize_handles(self, painter, rect):
        """绘制调整大小的手柄"""
        handle_size = 8
        half_handle = handle_size // 2

        # 四个角的手柄
        corners = [
            rect.topLeft(), rect.topRight(),
            rect.bottomLeft(), rect.bottomRight()
        ]

        # 四条边中点的手柄
        edges = [
            QPoint(rect.center().x(), rect.top()),
            QPoint(rect.center().x(), rect.bottom()),
            QPoint(rect.left(), rect.center().y()),
            QPoint(rect.right(), rect.center().y())
        ]

        painter.setPen(QPen(Qt.white, 1))
        painter.setBrush(QColor(255, 100, 100))

        for point in corners + edges:
            painter.drawRect(
                point.x() - half_handle,
                point.y() - half_handle,
                handle_size, handle_size
            )

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.handle_left_click(event)
        elif event.button() == Qt.MiddleButton:
            self.handle_middle_click(event)

    def handle_left_click(self, event):
        """处理鼠标左键点击"""
        if self.selected_region and self.selected_region.contains(event.pos()):
            # 检查是否点击了调整手柄
            edge = self.get_resize_edge(event.pos())
            if edge:
                self.start_resize(event.pos(), edge)
            else:
                # 开始移动选框
                self.is_moving = True
                self.move_offset = event.pos() - self.selected_region.topLeft()
        else:
            # 开始新选择
            self.is_selecting = True
            self.selection_start = event.pos()
            self.selection_end = event.pos()
            self.selected_region = None
            self.control_panel.hide()

        self.update()

    def handle_middle_click(self, event):
        """处理鼠标中键点击"""
        if self.selected_region and self.selected_region.contains(event.pos()):
            self.is_moving = True
            self.move_offset = event.pos() - self.selected_region.topLeft()
            self.setCursor(Qt.SizeAllCursor)

    def get_resize_edge(self, pos):
        """获取点击位置的调整边缘"""
        if not self.selected_region:
            return None

        rect = self.selected_region
        margin = 10

        # 检查四个角
        if (pos - rect.topLeft()).manhattanLength() < margin:
            return "topleft"
        elif (pos - rect.topRight()).manhattanLength() < margin:
            return "topright"
        elif (pos - rect.bottomLeft()).manhattanLength() < margin:
            return "bottomleft"
        elif (pos - rect.bottomRight()).manhattanLength() < margin:
            return "bottomright"

        # 检查四条边
        if abs(pos.y() - rect.top()) < margin and rect.left() <= pos.x() <= rect.right():
            return "top"
        elif abs(pos.y() - rect.bottom()) < margin and rect.left() <= pos.x() <= rect.right():
            return "bottom"
        elif abs(pos.x() - rect.left()) < margin and rect.top() <= pos.y() <= rect.bottom():
            return "left"
        elif abs(pos.x() - rect.right()) < margin and rect.top() <= pos.y() <= rect.bottom():
            return "right"

        return None

    def start_resize(self, pos, edge):
        """开始调整选框大小"""
        self.is_resizing = True
        self.resize_edge = edge
        self.resize_start_pos = pos
        self.resize_start_rect = QRect(self.selected_region)

    def mouseMoveEvent(self, event):
        if self.is_selecting:
            self.selection_end = event.pos()
        elif self.is_moving and self.selected_region:
            # 移动选框
            new_pos = event.pos() - self.move_offset
            self.selected_region.moveTo(new_pos)
            self.selection_start = self.selected_region.topLeft()
            self.selection_end = self.selected_region.bottomRight()
        elif self.is_resizing and self.selected_region:
            # 调整选框大小
            self.resize_selection(event.pos())

        self.update()

    def resize_selection(self, pos):
        """调整选框大小"""
        rect = QRect(self.resize_start_rect)
        delta = pos - self.resize_start_pos

        if self.resize_edge == "topleft":
            rect.setTopLeft(rect.topLeft() + delta)
        elif self.resize_edge == "topright":
            rect.setTopRight(rect.topRight() + delta)
        elif self.resize_edge == "bottomleft":
            rect.setBottomLeft(rect.bottomLeft() + delta)
        elif self.resize_edge == "bottomright":
            rect.setBottomRight(rect.bottomRight() + delta)
        elif self.resize_edge == "top":
            rect.setTop(rect.top() + delta.y())
        elif self.resize_edge == "bottom":
            rect.setBottom(rect.bottom() + delta.y())
        elif self.resize_edge == "left":
            rect.setLeft(rect.left() + delta.x())
        elif self.resize_edge == "right":
            rect.setRight(rect.right() + delta.x())

        # 确保选框最小尺寸
        if rect.width() > 10 and rect.height() > 10:
            self.selected_region = rect.normalized()
            self.selection_start = self.selected_region.topLeft()
            self.selection_end = self.selected_region.bottomRight()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton and self.is_selecting:
            self.is_selecting = False
            if self.selection_start and self.selection_end:
                rect = QRect(self.selection_start, self.selection_end).normalized()
                if rect.width() > 10 and rect.height() > 10:
                    self.selected_region = rect
                else:
                    self.selected_region = None

            # 退出全屏模式
            self.showNormal()

        elif event.button() == Qt.MiddleButton and self.is_moving:
            self.is_moving = False
            self.setCursor(Qt.CrossCursor)

        elif event.button() == Qt.LeftButton and self.is_resizing:
            self.is_resizing = False

        self.update()

    def save_screenshot(self):
        """保存截图"""
        if not self.selected_region:
            QMessageBox.warning(self, "警告", "请先选择截图区域")
            return

        filename, _ = QFileDialog.getSaveFileName(
            self, "保存截图", "1.png",
            "PNG 文件 (*.png);;JPEG 文件 (*.jpg *.jpeg);;所有文件 (*)")

        if filename:
            try:
                # 直接从内存中的QPixmap保存选择的区域
                rect = self.selected_region
                cropped = self.screenshot.copy(rect)

                # 根据文件扩展名确定保存格式
                if filename.lower().endswith(('.jpg', '.jpeg')):
                    cropped.save(filename, 'JPEG', quality=90)
                else:  # 默认保存为PNG
                    cropped.save(filename, 'PNG')

                QMessageBox.information(self, "成功", f"截图已保存到:\n{filename}")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存截图失败:\n{str(e)}")
            finally:
                self.close()

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Escape:
            if self.isFullScreen():
                self.showNormal()
            else:
                self.close()
